import json
from datetime import datetime
from enum import IntEnum
from typing import Optional, Callable, Any

from fastapi import WebSocket, WebSocketDisconnect
from fastapi.websockets import WebSocketState
from jose import JWTError
from pydantic import BaseModel, Field

from app.core.security import decode_token
from app.core.websocket import WebSocketConnectionManager
from app.utils.decorators.enum_description import enum_description


@enum_description(SYSTEM="系统消息", PUBLIC="群聊消息", PRIVATE="私聊消息", BYTES="二进制消息", ERROR="错误")
class MessageType(IntEnum):
    SYSTEM = 0
    PUBLIC = 1
    PRIVATE = 2
    BYTES = 3
    ERROR = 4

class Message(BaseModel):
    type: MessageType= Field(description="消息类型")
    recipient: str= Field(description="接收者")
    sender: str= Field(description="发送者")
    content: Any= Field(description="消息内容")
    timestamp: datetime = Field(default_factory=datetime.now,description="时间")


async def authenticate(token: str) -> Optional[str]:
    try:
        payload = decode_token(token)
        return payload.get("account")
    except JWTError:
        return None

class WebSocketController:
    def __init__(self, manager_factory: Callable[[], WebSocketConnectionManager]):
        self._manager_factory = manager_factory
        self._manager: Optional[WebSocketConnectionManager] = None

    @property
    def manager(self) -> WebSocketConnectionManager:
        if self._manager is None:
            self._manager = self._manager_factory()
        return self._manager

    async def handle_websocket_connection(self, websocket: WebSocket, token: str):
        username = await authenticate(token)
        if not username:
            await websocket.close(code=1008)
            return

        await self.manager.connect(websocket, username)

        try:
            await self.notify_user_joined(username)
            await self.handle_messages(websocket, username)
        except WebSocketDisconnect:
            await self.handle_disconnect(username)
        except Exception as e:
            # 检查连接是否仍然活跃
            if websocket.application_state == WebSocketState.CONNECTED:
                await self.handle_error(websocket, e)
            await self.handle_disconnect(username)
            raise

    async def notify_user_joined(self, username: str):
        try:
            if self.manager.is_connected(username):
                await self.manager.broadcast_json({
                    "type": "system",
                    "content": f"{username} 加入了聊天室",
                    "users": list(self.manager.active_connections.keys())
                })
        except Exception as e:
            print(f"广播用户加入消息失败: {str(e)}")

    async def handle_messages(self, websocket: WebSocket, username: str):
        try:
            while True:
                try:
                    data = await websocket.receive()

                    if data["type"] == "websocket.disconnect":
                        break

                    if data["type"] != "websocket.receive":
                        continue

                    if "text" in data:
                        await self.handle_text_data(data["text"], username, websocket)
                    elif "bytes" in data:
                        await self.handle_binary_data(data["bytes"], username, websocket)

                except WebSocketDisconnect:
                    break
                except Exception as e:
                    if websocket.application_state == WebSocketState.CONNECTED:
                        await self.handle_error(websocket, e)
                    break

        finally:
            await self.handle_disconnect(username)

    async def handle_text_data(self, text: str, username: str, websocket: WebSocket):
        try:
            data = json.loads(text)
            await self.handle_json_message(data, username, websocket)
        except json.JSONDecodeError:
            await self.handle_plain_text(text, username, websocket)

    async def handle_json_message(self, data: dict, username: str, websocket: WebSocket):
        message_type = data.get("type")

        if message_type == "private_message":
            await self.handle_private_message(data, username, websocket)
        else:
            await self.broadcast_message({
                "type": "public",
                "sender": username,
                "content": data.get("content"),
                "timestamp": datetime.now().isoformat()
            })

    async def handle_private_message(self, data: dict, sender: str, websocket: WebSocket):
        recipient = data.get("to")
        content = data.get("content")

        if not recipient or not content:
            return

        sent = await self.manager.send_private_json({
            "type": "private",
            "sender": sender,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }, recipient)

        if sent:
            await self.manager.send_json({
                "type": "private",
                "sender": sender,
                "to": recipient,
                "content": content,
                "timestamp": datetime.now().isoformat()
            }, websocket)
        else:
            await self.manager.send_json({
                "type": "error",
                "content": f"用户 {recipient} 不在线或不存在"
            }, websocket)

    async def handle_plain_text(self, text: str, username: str, websocket: WebSocket):
        if text.startswith("@"):
            await self.handle_private_command(text, username, websocket)
        else:
            await self.broadcast_message({
                "type": "public",
                "sender": username,
                "content": text,
                "timestamp": datetime.now().isoformat()
            })

    async def handle_private_command(self, text: str, sender: str, websocket: WebSocket):
        parts = text.split(" ", 1)
        if len(parts) != 2:
            return

        recipient = parts[0][1:]
        content = parts[1]

        sent = await self.manager.send_private_json({
            "type": "private",
            "sender": sender,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }, recipient)

        if sent:
            await self.manager.send_json({
                "type": "private",
                "to": recipient,
                "content": content
            }, websocket)
        else:
            await self.manager.send_text(f"用户 {recipient} 不在线或不存在", websocket)

    async def handle_binary_data(self, binary_data: bytes, username: str, websocket: WebSocket):
        await self.broadcast_message({
            "type": "binary_data",
            "sender": username,
            "size": len(binary_data),
            "timestamp": datetime.now().isoformat()
        })

        for connection in self.manager.active_connections.values():
            await connection.send_bytes(binary_data)

    async def broadcast_message(self, message: dict):
        await self.manager.broadcast_json(message)

    async def handle_disconnect(self, username: str):
        await self.manager.disconnect(username)
        await self.broadcast_message({
            "type": "system",
            "content": f"{username} 离开了聊天室",
            "users": list(self.manager.active_connections.keys())
        })

    async def handle_error(self, websocket: WebSocket, error: Exception):
        try:
            if websocket.application_state == WebSocketState.CONNECTED:
                await websocket.send_json({
                    "type": "error",
                    "content": f"发生错误: {str(error)}"
                })
        except Exception as e:
            print(f"发送错误消息失败: {str(e)}")